#include "Disambiguation.h"
#include <iostream>
#include <iomanip>
#include <string>
#include <queue>
#include <algorithm>
#include "FibHeap.h"

using namespace std;

Disambiguation::Disambiguation(SequenceGraph* seqGraph, ReadInputFacts& readFacts,
							   int pairSupport, int maxSeedCoverage, const char* outFile, const char* statisFile):
m_seqGraph(seqGraph), m_pairSupport(pairSupport), m_pairKSupport(pairSupport * (readFacts.getReadLength() - KmerVertex::getLength() + 1)), m_numOfReads(readFacts.numOfRead()),
m_readLength(readFacts.getReadLength()), m_insert(readFacts.getInsertionLength()), m_stddev(readFacts.getInsertionStdDev()),
m_maxSeedCoverage(maxSeedCoverage)
{
	m_readInSequence = new ReadInSequence[m_numOfReads+1]; // [0] is wasted
	memset(m_readInSequence, 0, (m_numOfReads+1) * sizeof(ReadInSequence));
	int arraySize = ((seqGraph->maxVertexID()) >> 3) + 1; // maxid / 8 bits + 1 makes enough space
	m_queued = new unsigned char[arraySize];
	memset(m_queued, 0, arraySize);
	m_seed = new unsigned char[arraySize];
	memset(m_seed, 0, arraySize);
	m_visited = new unsigned char[arraySize];
	memset(m_visited, 0, arraySize);
	m_outputed = new unsigned char[arraySize];
	memset(m_outputed, 0, arraySize);

	fcontig.open(outFile, ios::out);
	fstatis.open(statisFile);
}

void Disambiguation::allContigs()
{
	cout << "Start building all contigs." << endl;
	LOG(INFO) << "Start building all contigs.";
	m_longest = 0;
	m_numOfContigs = 0;
	seeding();
	//cout<<"seed Num: "<<m_seeds.size()<<endl;
	for (vector<GID_T>::iterator seed_iter = m_seeds.begin(); seed_iter != m_seeds.end(); ++seed_iter) {
		if (isQueued(*seed_iter)) {
			continue;
		}
		map<int, GID_T> skeleton; // coodinate, vid
		buildSkeleton(skeleton, *seed_iter);
		if (skeleton.size() < 2) {
			fcontig << ">CONTIG_" << ++m_numOfContigs <<endl;
			fcontig << m_seqGraph->getVertexByID(*seed_iter)->decompressVertex(*m_seqGraph, shared_buffer);
			setOutputed(*seed_iter);
			fcontig << endl;
			int contig_len = m_seqGraph->getVertexByID(*seed_iter)->size();
			double contig_cov = m_seqGraph->getVertexByID(*seed_iter)->coverage();
			fstatis << ">CONTIG_" << m_numOfContigs << ' ' << contig_len << ' ' << setprecision(4) << contig_cov << endl;
			m_longest = max(m_longest, contig_len);
			continue;
		}
		int contig_len = 0;
		double contig_obs = 0;
		fcontig << ">CONTIG_" <<++m_numOfContigs <<endl;
		vector<GID_T> contig; // ID or 0 following how many N to fill
		map<int, GID_T>::iterator liter = skeleton.begin();
		contig.push_back(liter->second);
		map<int, GID_T>::iterator riter = skeleton.begin();
		for (++riter; riter != skeleton.end(); ++riter) {
			deque<GID_T> bridge;
			pave(liter->second, riter->second, riter->first - liter->first, bridge);
			if (bridge.size() > 3 && bridge[1] == 0 && bridge[2] < 0) {
				continue; // skip too close
			} else {
				contig.insert(contig.end(), ++(bridge.begin()), bridge.end());
				liter = riter;
			}
		} // for (++riter; riter != skeleton.end(); ++liter,++riter)
		int toSkip = 0;
		for (int i = 0; i < contig.size(); ++i) {
			if (contig[i] != 0) {
				m_seqGraph->getVertexByID(contig[i])->decompressVertex(*m_seqGraph, shared_buffer, toSkip);
				int additional = m_seqGraph->getVertexByID(contig[i])->size() - toSkip;
				contig_len += additional;
				contig_obs += additional * m_seqGraph->getVertexByID(contig[i])->coverage();
				toSkip = KmerVertex::getLength() - 1;
				fcontig << shared_buffer;
				setOutputed(contig[i]);
			} // if (contig[i] != 0)
			else {
				int numOfUnsure = contig[++i];
				if (numOfUnsure > 0) {
					contig_len += numOfUnsure;
				}
				for (int j = 0; j < numOfUnsure; ++j) {
					fcontig << 'N';
				}
				toSkip = numOfUnsure > 0 ? 0 : -numOfUnsure;
			}
		} // for (int i = 0; i < contig.size(); ++i)
		fcontig << endl;
		fstatis << ">CONTIG_" << m_numOfContigs << ' ' << contig_len << ' ' << setprecision(4) << (contig_obs / contig_len) << endl;
		m_longest = max(m_longest, contig_len);
	} // for (vector<GID_T>::iterator seed_iter = m_seeds.begin(); seed_iter != m_seeds.end(); ++seed_iter)
	
	m_seqGraph->resetCursor();
	while (m_seqGraph->hasNext())
	{
		SequenceVertex* curSeqV = m_seqGraph->next();
		if (!isOutputed(curSeqV->getVertexID())) {
			deque<GID_T> contig;
			tryExtend(curSeqV, contig);
			int toSkip = 0;
			int contig_len = 0;
			fcontig << ">CONTIG_" <<++m_numOfContigs <<endl;
			for (deque<GID_T>::iterator it = contig.begin(); it != contig.end(); ++it) {
				m_seqGraph->getVertexByID(*it)->decompressVertex(*m_seqGraph, shared_buffer, toSkip);
				contig_len += m_seqGraph->getVertexByID(*it)->size() - toSkip;
				toSkip = KmerVertex::getLength() - 1;
				fcontig << shared_buffer;
				setOutputed(*it);
			} // for (deque<GID_T>::iterator it = contig.begin(); it != contig.end(); ++it)
			fcontig << endl;
			m_longest = max(m_longest, contig_len);
		} // if (!isOutputed(curSeqV->getVertexID()))
	}
	cout << "Finish building all contigs." << endl;
	LOG(INFO) << "Finish building all contigs.";
}

void Disambiguation::seeding()
{
	m_seqGraph->resetCursor();
	while (m_seqGraph->hasNext()) {
		SequenceVertex* v = m_seqGraph->next();
		SequenceVertex::ReadsOffsetAndConfidence& readoac = v->getReadsInfo();
		for (SequenceVertex::ConstReadsIterator iter = readoac.begin(); iter != readoac.end(); ++iter) {
			GID_T readid = iter->first;
			int offset = iter->second.first;
			GID_T vertexid = v->getVertexID();
			if (readid < 0) {
				readid = -readid;
				offset = v->size() - m_readLength - offset;
				vertexid = -vertexid;
			} // if (readid < 0)
			m_readInSequence[readid].seqVertexID = vertexid;
			m_readInSequence[readid].offset = offset;
			m_readInSequence[readid].confidence = iter->second.second;
		}
		if (v->size() >= m_pairSupport + m_readLength - 1 && v->coverage() <= m_maxSeedCoverage) { // criteria could change
			m_seeds.push_back(v->getVertexID());
			setSeed(v->getVertexID());
			//m_aux.insert(make_pair(v->getVertexID(), make_pair(0, false)));
		}
	} // while (m_seqGraph->hasNext())
	sort(m_seeds.begin(), m_seeds.end(), size_compare(m_seqGraph));
}

void Disambiguation::buildSkeleton(map<int, GID_T>& skeleton, GID_T init_seed)
{
	queue<pair<GID_T, int> > seedQ; // id, coordinate
	skeleton.insert(make_pair(0, init_seed));
	seedQ.push(make_pair(init_seed, 0));
	setQueued(init_seed);
	while (!seedQ.empty()) {
		const pair<GID_T, int>& seedCood = seedQ.front();
		SequenceVertex* vertex = m_seqGraph->getVertexByID(seedCood.first);
		const SequenceVertex::ReadsOffsetAndConfidence& readInfo = vertex->getReadsInfo();
		map<GID_T, pair<vector<int>, int> > pairedVertexInfo; // id, <estimatedCoodinate list, support>
		// test each read in vertex to find neighbor paired vertices
		for (SequenceVertex::ConstReadsIterator iter = readInfo.begin(); iter != readInfo.end(); ++iter) {
			const PairedSequenceInfo& paired = getPairedSequence(iter->first, iter->second.first, seedCood.second);
			if (!paired.seqVertexID) { // paired sequence node not found
				continue;
			}
			map<GID_T, pair<vector<int>, int> >::iterator findIter = pairedVertexInfo.find(paired.seqVertexID);
			if (findIter == pairedVertexInfo.end()) {
				findIter = pairedVertexInfo.insert(make_pair(paired.seqVertexID, make_pair(vector<int>(), 0))).first;
			}
			findIter->second.first.push_back(paired.coordinate);
			findIter->second.second += min(paired.confidence, iter->second.second);
		} // for (SequenceVertex::ConstReadsIterator iter = readInfo.begin(); iter != readInfo.end(); ++iter)
		// Add trusted into skeleton
		for (map<GID_T, pair<vector<int>, int> >::iterator iter = pairedVertexInfo.begin();
			iter != pairedVertexInfo.end(); ++iter) {
				if (iter->second.second < m_pairKSupport) {
					continue; // not trusted
				}
				vector<int>& coordList = iter->second.first;
				sort(coordList.begin(), coordList.end());
				int median = coordList[coordList.size() >> 1];
				// avoid repeated insertion caused by a little deviate
				int lowest = median - m_stddev;
				int highest = median + m_stddev;
				map<int, GID_T>::iterator lowIter = skeleton.lower_bound(lowest);
				map<int, GID_T>::iterator highIter = skeleton.upper_bound(highest);
				for (; lowIter != highIter; ++lowIter) {
					if (lowIter->second == iter->first) {
						break;
					}
				} // for (; lowIter != highIter; ++lowIter)
				if (lowIter == highIter) { // not repeat, so insert
					skeleton.insert(make_pair(median, iter->first));
					if (isSeed(iter->first) && !isQueued(iter->first)) {
						seedQ.push(make_pair(iter->first, median));
						setQueued(iter->first);
					}
				}
		} // for (map<GID_T, pair<vertex<int>, int> >::iterator iter = pairedVertexInfo.begin();iter != pairedVertexInfo.end(); ++iter)
		seedQ.pop();
	} // while (!seedQ.empty())
}

Disambiguation::PairedSequenceInfo Disambiguation::getPairedSequence(GID_T readId, int offset, int coordinate)
{
	if (readId > 0) {
		if (readId % 2) { // odd
			ReadInSequence& pairedSequence = m_readInSequence[readId + 1];
			if (pairedSequence.seqVertexID == 0) {
				return NoPairedSequenceInfo;
			}
			return PairedSequenceInfo(pairedSequence.seqVertexID , 
				coordinate + offset + m_insert - m_readLength - pairedSequence.offset, pairedSequence.confidence);
		} // if (readId % 2)
		else {
			ReadInSequence& pairedSequence = m_readInSequence[readId - 1];
			if (pairedSequence.seqVertexID == 0) {
				return NoPairedSequenceInfo;
			}
			return PairedSequenceInfo(pairedSequence.seqVertexID,
				coordinate + offset + m_readLength - m_insert - pairedSequence.offset, pairedSequence.confidence);
		}
	} // if (readId > 0)
	if (readId % 2) { // odd
		ReadInSequence& pairedSequence = m_readInSequence[-(readId - 1)];
		if (pairedSequence.seqVertexID == 0) {
			return NoPairedSequenceInfo;
		}
		SequenceVertex* pv = m_seqGraph->getVertexByID(-pairedSequence.seqVertexID);
		return PairedSequenceInfo(-pairedSequence.seqVertexID , 
			coordinate + offset + m_readLength - m_insert - (pv->size() - m_readLength - pairedSequence.offset),
			pairedSequence.confidence);
	} // if (readId % 2)
	else {
		ReadInSequence& pairedSequence = m_readInSequence[-(readId + 1)];
		if (pairedSequence.seqVertexID == 0) {
			return NoPairedSequenceInfo;
		}
		SequenceVertex* pv = m_seqGraph->getVertexByID(-(pairedSequence.seqVertexID));
		return PairedSequenceInfo(-(pairedSequence.seqVertexID),
			coordinate + offset + m_insert - m_readLength - (pv->size() - m_readLength - pairedSequence.offset),
			pairedSequence.confidence);
	}
}

void Disambiguation::pave(GID_T startNodeID, GID_T endNodeID, int meanLength, deque<GID_T>& bridge)
{
	int minLength = meanLength - 2 * m_stddev;
	int maxLength = meanLength + 2 * m_stddev;
	SequenceVertex* startNode = m_seqGraph->getVertexByID(startNodeID);
	SequenceVertex* endNode = m_seqGraph->getVertexByID(endNodeID);

	// first try the shortest path
	FibHeap heap(PathNode::cmp_func);
	PathNode root(startNodeID);
	vector<AbstractHeapNode*> allocatedPathNodes;
	clearVisited();
	heap.insert(&root);
	setVisited(startNodeID);
	bool everMet = false;
	while (!heap.isEmpty()) {
		PathNode* curPN = static_cast<PathNode*>(heap.extractMin());
		SequenceVertex* curSV = m_seqGraph->getVertexByID(curPN->seqVertexID);
		const vector<SequenceEdge>& succs = m_seqGraph->getSuccessors(curSV);
		for (int i = 0; i < succs.size(); ++i) {
			int nextDist = curPN->dist + curSV->size() - KmerVertex::getLength() + 1;
			if (!isVisited(succs[i].end->getVertexID()) && nextDist <= maxLength) {
				PathNode* newPN = new PathNode(succs[i].end->getVertexID(), curPN);
				allocatedPathNodes.push_back(newPN);
				newPN->dist = nextDist;
				if (newPN->seqVertexID == endNodeID) {
					everMet = true;
					if (newPN->dist > minLength && newPN->dist < maxLength) {
						bridge.push_front(newPN->seqVertexID);
						while ((newPN = newPN->parent) != NULL) {
							bridge.push_front(newPN->seqVertexID);
						}
						cleanAllocated(allocatedPathNodes);
						return;
					} // if (newPN->dist > minLength && newPN->dist < maxLength)
				} // if (newPN->seqVertexID == endNodeID)
				heap.insert(newPN);
				setVisited(newPN->seqVertexID);
			} // if (!isVisited(succs[i]->getVertexID()) && nextDist <= maxLength)
		} // for (int i = 0; i < succs.size(); ++i)
	} // while (!heap.isEmpty())
	cleanAllocated(allocatedPathNodes);

	// if the shortest path is not the answer.

	if (!everMet) { // target node is not in range, strange ...
		bridge.push_back(startNodeID);
		bridge.push_back(0);
		bridge.push_back(meanLength - startNode->size());
		bridge.push_back(endNodeID);
		return;
	} //if (!everMet)

	//generate candidates before brute force search
	set<GID_T> candidates;
	generateCandidates(candidates, startNodeID, endNodeID, maxLength);
	bruteForcePave(startNodeID, endNodeID, meanLength, bridge, candidates);
}

void Disambiguation::generateCandidates(set<GID_T>& candidates, GID_T startNodeID, GID_T endNodeID, int maxLength)
{
	int limit = 1000000;
	FibHeap heap(CandidateNode::cmp_func);
	map<GID_T, CandidateNode> toCandidate;
	map<GID_T, CandidateNode>::iterator insertIter = toCandidate.insert( make_pair(startNodeID, CandidateNode(startNodeID)) ).first;
	heap.insert(&(insertIter->second));
	while (!heap.isEmpty()) {
		if (--limit < 0) { break; }
		CandidateNode* curCN = static_cast<CandidateNode*>(heap.extractMin());
		SequenceVertex* curSV = m_seqGraph->getVertexByID(curCN->seqVertexID);
		const vector<SequenceEdge>& succs = m_seqGraph->getSuccessors(curSV);
		for (int i = 0; i < succs.size(); ++i) {
			int nextDist = curCN->dist + curSV->size() - KmerVertex::getLength() + 1;
			if (nextDist <= maxLength) {
				GID_T succID = succs[i].end->getVertexID();
				map<GID_T, CandidateNode>::iterator findIter = toCandidate.find(succID);
				if (findIter == toCandidate.end()) {
					findIter = toCandidate.insert( make_pair(succID, CandidateNode(succID)) ).first;
				}
				findIter->second.dist = nextDist;
				findIter->second.parents.push_back(curCN);
				heap.insert(&(findIter->second));
			} // if (nextDist <= maxLength)
		} // for (int i = 0; i < succs.size(); ++i)
	} // while (!heap.isEmpty())
	map<GID_T, CandidateNode>::iterator findIter = toCandidate.find(endNodeID);
	if (findIter == toCandidate.end()) { return; }
	CandidateNode* curNode = &(findIter->second);
	queue<CandidateNode*> qc;
	qc.push(curNode);
	candidates.insert(curNode->seqVertexID);
	while (!qc.empty()) {
		curNode = qc.front();
		for (int i = 0; i < curNode->parents.size(); ++i) {
			CandidateNode* parent = curNode->parents[i];
			if (candidates.insert(parent->seqVertexID).second) {
				qc.push(parent);
			}
		} // for (int i = 0; i < curNode->parents.size(); ++i)
		qc.pop();
	}
}

void Disambiguation::bruteForcePave(GID_T startNodeID, GID_T endNodeID, int meanLength, deque<GID_T>& bridge, set<GID_T>& candidates)
{
	int limit = 1000000;

	int minLength = meanLength - 2 * m_stddev;
	int maxLength = meanLength + 2 * m_stddev;
	SequenceVertex* startNode = m_seqGraph->getVertexByID(startNodeID);
	SequenceVertex* endNode = m_seqGraph->getVertexByID(endNodeID);

	FibHeap heap(PathNode::cmp_func);
	PathNode root(startNodeID);
	vector<AbstractHeapNode*> allocatedPathNodes;
	heap.insert(&root);
	while (!heap.isEmpty()) {
		if (--limit < 0) { break; }

		PathNode* curPN = static_cast<PathNode*>(heap.extractMin());
		SequenceVertex* curSV = m_seqGraph->getVertexByID(curPN->seqVertexID);
		const vector<SequenceEdge>& succs = m_seqGraph->getSuccessors(curSV);
		for (int i = 0; i < succs.size(); ++i) {
			int nextDist = curPN->dist + curSV->size() - KmerVertex::getLength() + 1;
			if (candidates.find(succs[i].end->getVertexID()) != candidates.end() && nextDist <= maxLength) {
				PathNode* newPN = new PathNode(succs[i].end->getVertexID(), curPN);
				allocatedPathNodes.push_back(newPN);
				newPN->dist = nextDist;
				if (newPN->seqVertexID == endNodeID &&
					newPN->dist > minLength && newPN->dist < maxLength) {
					bridge.push_front(newPN->seqVertexID);
					while ((newPN = newPN->parent) != NULL) {
						bridge.push_front(newPN->seqVertexID);
					}
					cleanAllocated(allocatedPathNodes);
					return;
				} // if (newPN->seqVertexID == endNodeID)
				heap.insert(newPN);
			} // if (!isVisited(succs[i]->getVertexID()) && nextDist <= maxLength)
		} // for (int i = 0; i < succs.size(); ++i)
	} // while (!heap.isEmpty())
	cleanAllocated(allocatedPathNodes);

	// all attempts have failed fill it with N
	bridge.push_back(startNodeID);
	bridge.push_back(0);
	bridge.push_back(meanLength - startNode->size());
	bridge.push_back(endNodeID);
	return;
}

void Disambiguation::tryExtend(SequenceVertex* midVertex, deque<GID_T>& contig) {
	clearVisited();
	contig.push_back(midVertex->getVertexID());
	setVisited(midVertex->getVertexID());
	const vector<SequenceEdge>* succs = &(midVertex->getSuccessors());
	while (succs && !succs->empty()) {
		SequenceVertex* noutSucc = NULL;
		int scnt = 0;
		for (vector<SequenceEdge>::const_iterator it = succs->begin(); it != succs->end(); ++ it) {
			SequenceVertex* neigh = it->end;
			if (!isOutputed(neigh->getVertexID()) && !isVisited(neigh->getVertexID())) {
				noutSucc = neigh;
				++scnt;
			} // if (!isOutputed(neigh->getVertexID()) && !isVisited(neigh->getVertexID()))
		} // for (vector<SequenceEdge>::const_iterator it = succs->begin(); it != succs->end(); ++ it) 
		if (scnt != 1) {
			succs = NULL;
		} else {
			contig.push_back(noutSucc->getVertexID());
			setVisited(noutSucc->getVertexID());
			succs = &(noutSucc->getSuccessors());
		}
	} // while (succs) 
	const vector<SequenceEdge>* preds = &(midVertex->getTwinVertex()->getSuccessors());
	while (preds && !preds->empty()) {
		SequenceVertex* noutPred = NULL;
		int pcnt = 0;
		for (vector<SequenceEdge>::const_iterator it = preds->begin(); it != preds->end(); ++it) {
			SequenceVertex* neigh = it->end->getTwinVertex();
			if (!isOutputed(neigh->getVertexID()) && !isVisited(neigh->getVertexID())) {
				noutPred = neigh;
				++pcnt;
			} // if (!isOutputed(neigh->getVertexID()) && !isVisited(neigh->getVertexID()))
		} // for (vector<SequenceEdge>::const_iterator it = preds->begin(); it != preds->end(); ++ it) 
		if (pcnt != 1) {
			preds = NULL;
		} else {
			contig.push_front(noutPred->getVertexID());
			setVisited(noutPred->getVertexID());
			preds = &(noutPred->getPredecessors());
		}
	} // while (preds) 
}

void Disambiguation::cleanAllocated(vector<AbstractHeapNode*>& allocated) {
	for (GUINT32 i = 0; i < allocated.size(); ++i) {
		delete allocated[i];
	}
	allocated.clear();
}

Disambiguation::~Disambiguation()
{
	cerr <<"~Disambiguation" << endl;
	fcontig.close();
	fstatis.close();
	delete[] m_visited;
	delete[] m_queued;
	delete[] m_seed;
	delete[] m_outputed;
	cerr <<"exit ~Disambiguation" << endl;
}
